package com.eshipenet.shipnet.edgeboxextbuilder
import org.apache.flink.api.common.state.MapStateDescriptor
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction
import org.apache.flink.util.Collector

class EdgeBoxActionProcessFunction(desc: MapStateDescriptor[String, EdgeBoxModel])
    extends BroadcastProcessFunction[EdgeBoxAction, EdgeBoxModelAction, EdgeBoxExtAction] {

  override def processElement(
      value: EdgeBoxAction,
      ctx: BroadcastProcessFunction[EdgeBoxAction, EdgeBoxModelAction, EdgeBoxExtAction]#ReadOnlyContext,
      out: Collector[EdgeBoxExtAction]): Unit = {
    val models = ctx.getBroadcastState(desc)
    val act = value match {
      case insert @ EdgeBoxInsert(id, after, _) =>
        val model = Option(models.get(after.modelId))
        EdgeBoxExtInsert(id, insert, model)
      case update @ EdgeBoxUpdate(id, before, after, _) =>
        val model = Option(models.get(after.modelId))
        EdgeBoxExtUpdate(id, update, model)
      case delete @ EdgeBoxDelete(id, _, _) =>
        EdgeBoxExtDelete(id, delete)
    }
    out.collect(act)
  }

  override def processBroadcastElement(
      value: EdgeBoxModelAction,
      ctx: BroadcastProcessFunction[EdgeBoxAction, EdgeBoxModelAction, EdgeBoxExtAction]#Context,
      out: Collector[EdgeBoxExtAction]): Unit = {
    val models = ctx.getBroadcastState(desc)
    value match {
      case EdgeBoxModelInsert(id, after) =>
        models.put(id, after)
      case EdgeBoxModelUpdate(id, before, after) =>
        models.put(id, after)
      case EdgeBoxModelDelete(id, _) =>
        models.remove(id)
    }
  }

}
